package com.gcp.ai.service.impl;

import com.gcp.ai.service.SseService;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SseServiceImpl implements SseService {

    private static final ConcurrentHashMap<String, SseEmitter> sseCache = new ConcurrentHashMap<>();

    @Override
    public SseEmitter createSseConnect(String clientId) {
        SseEmitter sseEmitter = new SseEmitter(0L);
        sseCache.put(clientId, sseEmitter);
        sseEmitter.onCompletion(() -> {
            sseCache.get(clientId).complete();
            sseCache.remove(clientId);
        });
        sseEmitter.onTimeout(() -> {
            sseCache.get(clientId).complete();
            sseCache.remove(clientId);
        });
        sseEmitter.onError((e) -> {
            sseCache.get(clientId).complete();
            sseCache.remove(clientId);
        });
        return sseEmitter;
    }

    @Override
    public void sendMessage(String clientId,String msg) {
        if(sseCache.get(clientId) == null){
            return;
        }
        try {
            SseEmitter.SseEventBuilder message = SseEmitter.event()
                    .id("1")
                    .data(msg);
            sseCache.get(clientId).send(message);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void complete(String clientId) {
        if(sseCache.get(clientId) == null){
            return;
        }
        try {
            sseCache.get(clientId).complete();
            sseCache.remove(clientId);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
